{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from py2neo import Graph\n",
    "\n",
    "class AnswerSearcher:\n",
    "    def __init__(self):\n",
    "        # self.g = Graph(\"http://localhost:7474\", username=\"neo4j\", password=\"123456\")\n",
    "        self.g = Graph(\"http://localhost:7474\", auth=(\"neo4j\", \"123456\"))\n",
    "        self.num_limit = 20\n",
    "\n",
    "    '''执行cypher查询，并返回相应结果'''\n",
    "    def search_main(self, sqls):\n",
    "        final_answers = []\n",
    "        for sql_ in sqls:\n",
    "            question_type = sql_['question_type']\n",
    "            queries = sql_['sql']\n",
    "            answers = []\n",
    "            for query in queries:\n",
    "                ress = self.g.run(query).data()\n",
    "                answers += ress\n",
    "            final_answer = self.answer_prettify(question_type, answers)\n",
    "            if final_answer:\n",
    "                final_answers.append(final_answer)\n",
    "        return final_answers\n",
    "\n",
    "    '''根据对应的qustion_type，调用相应的回复模板'''\n",
    "    def answer_prettify(self, question_type, answers):\n",
    "        final_answer = []\n",
    "        if not answers:\n",
    "            return ''\n",
    "        # 疾病-症状\n",
    "        if question_type == 'disease_symptom':\n",
    "            desc = [i['n.name'] for i in answers]\n",
    "            subject = answers[0]['m.name']\n",
    "            final_answer = '{0}的症状包括：{1}'.format(subject, '；'.join(list(set(desc))[:self.num_limit]))\n",
    "        # 症状-疾病\n",
    "        elif question_type == 'symptom_disease':\n",
    "            desc = [i['m.name'] for i in answers]\n",
    "            subject = answers[0]['n.name']\n",
    "            final_answer = '症状{0}可能染上的疾病有：{1}'.format(subject, '；'.join(list(set(desc))[:self.num_limit]))\n",
    "\n",
    "        return final_answer\n",
    "\n",
    "searcher = AnswerSearcher()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class QuestionPaser:\n",
    "\n",
    "    '''构建实体节点'''\n",
    "    def build_entitydict(self, args):\n",
    "        entity_dict = {}\n",
    "        for arg, types in args.items():\n",
    "            for type in types:\n",
    "                if type not in entity_dict:\n",
    "                    entity_dict[type] = [arg]\n",
    "                else:\n",
    "                    entity_dict[type].append(arg)\n",
    "\n",
    "        return entity_dict\n",
    "\n",
    "    '''解析主函数'''\n",
    "    def parser_main(self, res_classify):\n",
    "        args = res_classify['args']\n",
    "        entity_dict = self.build_entitydict(args)\n",
    "        question_types = res_classify['question_types']\n",
    "        sqls = []\n",
    "        for question_type in question_types:\n",
    "            sql_ = {}\n",
    "            sql_['question_type'] = question_type\n",
    "            sql = []\n",
    "            if question_type == 'disease_symptom': # 疾病-症状\n",
    "                sql = self.sql_transfer(question_type, entity_dict.get('disease'))\n",
    "\n",
    "            elif question_type == 'symptom_disease': # 症状-疾病\n",
    "                sql = self.sql_transfer(question_type, entity_dict.get('symptom'))\n",
    "\n",
    "\n",
    "            if sql:\n",
    "                sql_['sql'] = sql\n",
    "\n",
    "                sqls.append(sql_)\n",
    "\n",
    "        return sqls\n",
    "\n",
    "    '''针对不同的问题，分开进行处理'''\n",
    "    def sql_transfer(self, question_type, entities):\n",
    "        if not entities:\n",
    "            return []\n",
    "\n",
    "        # 查询语句\n",
    "        sql = []\n",
    "        # 查询疾病的原因\n",
    "        if question_type == 'disease_cause':\n",
    "            sql = [\"MATCH (m:Disease) where m.name = '{0}' return m.name, m.cause\".format(i) for i in entities]\n",
    "\n",
    "        # 查询疾病的防御措施\n",
    "        elif question_type == 'disease_prevent':\n",
    "            sql = [\"MATCH (m:Disease) where m.name = '{0}' return m.name, m.prevent\".format(i) for i in entities]\n",
    "\n",
    "        # 查询疾病的持续时间\n",
    "        elif question_type == 'disease_lasttime':\n",
    "            sql = [\"MATCH (m:Disease) where m.name = '{0}' return m.name, m.cure_lasttime\".format(i) for i in entities]\n",
    "\n",
    "        # 查询疾病的治愈概率\n",
    "        elif question_type == 'disease_cureprob':\n",
    "            sql = [\"MATCH (m:Disease) where m.name = '{0}' return m.name, m.cured_prob\".format(i) for i in entities]\n",
    "\n",
    "        # 查询疾病的治疗方式\n",
    "        elif question_type == 'disease_cureway':\n",
    "            sql = [\"MATCH (m:Disease) where m.name = '{0}' return m.name, m.cure_way\".format(i) for i in entities]\n",
    "\n",
    "        # 查询疾病的易发人群\n",
    "        elif question_type == 'disease_easyget':\n",
    "            sql = [\"MATCH (m:Disease) where m.name = '{0}' return m.name, m.easy_get\".format(i) for i in entities]\n",
    "\n",
    "        # 查询疾病的相关介绍\n",
    "        elif question_type == 'disease_desc':\n",
    "            sql = [\"MATCH (m:Disease) where m.name = '{0}' return m.name, m.desc\".format(i) for i in entities]\n",
    "\n",
    "        # 查询疾病有哪些症状\n",
    "        elif question_type == 'disease_symptom':\n",
    "            sql = [\"MATCH (m:Disease)-[r:has_symptom]->(n:Symptom) where m.name = '{0}' return m.name, r.name, n.name\".format(i) for i in entities]\n",
    "\n",
    "        # 查询症状会导致哪些疾病\n",
    "        elif question_type == 'symptom_disease':\n",
    "            sql = [\"MATCH (m:Disease)-[r:has_symptom]->(n:Symptom) where n.name = '{0}' return m.name, r.name, n.name\".format(i) for i in entities]\n",
    "\n",
    "        # 查询疾病的并发症\n",
    "        elif question_type == 'disease_acompany':\n",
    "            sql1 = [\"MATCH (m:Disease)-[r:acompany_with]->(n:Disease) where m.name = '{0}' return m.name, r.name, n.name\".format(i) for i in entities]\n",
    "            sql2 = [\"MATCH (m:Disease)-[r:acompany_with]->(n:Disease) where n.name = '{0}' return m.name, r.name, n.name\".format(i) for i in entities]\n",
    "            sql = sql1 + sql2\n",
    "        # 查询疾病的忌口\n",
    "        elif question_type == 'disease_not_food':\n",
    "            sql = [\"MATCH (m:Disease)-[r:no_eat]->(n:Food) where m.name = '{0}' return m.name, r.name, n.name\".format(i) for i in entities]\n",
    "\n",
    "        # 查询疾病建议吃的东西\n",
    "        elif question_type == 'disease_do_food':\n",
    "            sql1 = [\"MATCH (m:Disease)-[r:do_eat]->(n:Food) where m.name = '{0}' return m.name, r.name, n.name\".format(i) for i in entities]\n",
    "            sql2 = [\"MATCH (m:Disease)-[r:recommand_eat]->(n:Food) where m.name = '{0}' return m.name, r.name, n.name\".format(i) for i in entities]\n",
    "            sql = sql1 + sql2\n",
    "\n",
    "        # 已知忌口查疾病\n",
    "        elif question_type == 'food_not_disease':\n",
    "            sql = [\"MATCH (m:Disease)-[r:no_eat]->(n:Food) where n.name = '{0}' return m.name, r.name, n.name\".format(i) for i in entities]\n",
    "\n",
    "        # 已知推荐查疾病\n",
    "        elif question_type == 'food_do_disease':\n",
    "            sql1 = [\"MATCH (m:Disease)-[r:do_eat]->(n:Food) where n.name = '{0}' return m.name, r.name, n.name\".format(i) for i in entities]\n",
    "            sql2 = [\"MATCH (m:Disease)-[r:recommand_eat]->(n:Food) where n.name = '{0}' return m.name, r.name, n.name\".format(i) for i in entities]\n",
    "            sql = sql1 + sql2\n",
    "\n",
    "        # 查询疾病常用药品－药品别名记得扩充\n",
    "        elif question_type == 'disease_drug':\n",
    "            sql1 = [\"MATCH (m:Disease)-[r:common_drug]->(n:Drug) where m.name = '{0}' return m.name, r.name, n.name\".format(i) for i in entities]\n",
    "            sql2 = [\"MATCH (m:Disease)-[r:recommand_drug]->(n:Drug) where m.name = '{0}' return m.name, r.name, n.name\".format(i) for i in entities]\n",
    "            sql = sql1 + sql2\n",
    "\n",
    "        # 已知药品查询能够治疗的疾病\n",
    "        elif question_type == 'drug_disease':\n",
    "            sql1 = [\"MATCH (m:Disease)-[r:common_drug]->(n:Drug) where n.name = '{0}' return m.name, r.name, n.name\".format(i) for i in entities]\n",
    "            sql2 = [\"MATCH (m:Disease)-[r:recommand_drug]->(n:Drug) where n.name = '{0}' return m.name, r.name, n.name\".format(i) for i in entities]\n",
    "            sql = sql1 + sql2\n",
    "        # 查询疾病应该进行的检查\n",
    "        elif question_type == 'disease_check':\n",
    "            sql = [\"MATCH (m:Disease)-[r:need_check]->(n:Check) where m.name = '{0}' return m.name, r.name, n.name\".format(i) for i in entities]\n",
    "\n",
    "        # 已知检查查询疾病\n",
    "        elif question_type == 'check_disease':\n",
    "            sql = [\"MATCH (m:Disease)-[r:need_check]->(n:Check) where n.name = '{0}' return m.name, r.name, n.name\".format(i) for i in entities]\n",
    "\n",
    "        return sql"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
